#include "sv2_autofocus_det.h"
#include <cmath>
#include <opencv2/calib3d.hpp>
#include <opencv2/dnn.hpp>
#include <fstream>


namespace sv2 {

using namespace cv;
using namespace cv::dnn;


void AutoFocusDetCppNode::run()
{
  while (this->is_running())
  {
    nlohmann::json image_json;
    nlohmann::json tgts_json;

    if (this->_locked && this->_lock_rois.size() > 0)
    {
      // 主线程阻塞，等待队里中出现数据
      std::unique_lock<std::mutex> lock(this->_roi_queue_mtx);
      this->_roi_cv.wait(lock, [this] { return !this->_roi_queue.empty(); });

      image_json = this->_roi_queue.front();
      this->_roi_queue.pop();
      tgts_json = image_json["spirecv_msgs::2DTargets"];

      if (tgts_json.contains("rois") && tgts_json["rois"].size() > 0)
      {
        ;
      }
      else
      {
        continue;
      }
    }
    else
    {
      // 主线程阻塞，等待队里中出现数据
      std::unique_lock<std::mutex> lock(this->_full_queue_mtx);
      this->_full_cv.wait(lock, [this] { return !this->_full_queue.empty(); });

      image_json = this->_full_queue.front();
      this->_full_queue.pop();
      tgts_json = image_json["spirecv_msgs::2DTargets"];

      if (tgts_json.contains("rois") && tgts_json["rois"].size() > 0)
      {
        continue;
      }
    }


    bool found = false;
    bool outputing = false;
    int img_h = tgts_json["height"], img_w = tgts_json["width"];

    nlohmann::json _min_dist_tgt;
    nlohmann::json _min_dist_roi;
    float _min_dist_cx;
    float _min_dist_cy;
    float _min_dist = 1e8;
    int _min_dist_cnt = 0;

    if (!this->_locked)
    {
      nlohmann::json max_conf_inst;
      float max_conf = 0.f;
      for (int i=0; i<(int)tgts_json["targets"].size(); i++)
      {
        if (tgts_json["targets"][i]["score"] > max_conf)
        {
          max_conf = tgts_json["targets"][i]["score"];
          max_conf_inst = tgts_json["targets"][i];
        }
      }
      if (max_conf > 0.01f)
      {
        tgts_json["targets"].clear();
        tgts_json["targets"].push_back(max_conf_inst);
      }
    }

  
    for (int i=0; i<(int)tgts_json["targets"].size(); i++)
    {
      std::vector<float> b; 
      for (int j=0; j<tgts_json["targets"][i]["bbox"].size(); j++)
        b.push_back(tgts_json["targets"][i]["bbox"][j]);

      bool is_inst_in = true;
      float wh = std::max(b[2], b[3]) * this->_lock_scale_init;
      if (this->_lock_scale_count > 15)
        wh = std::max(b[2], b[3]) * this->_lock_scale;
      if (wh > std::min(img_h, img_w))
        wh = std::min(img_h, img_w);
      if (wh < 64)
        wh = 64;
      float cx = b[0] + b[2] / 2.f;
      float cy = b[1] + b[3] / 2.f;
      if (this->_locked)
      {
        if (cx < this->_lock_rois[0][0] || cx > this->_lock_rois[0][2])
          is_inst_in = false;
        if (cy < this->_lock_rois[0][1] || cy > this->_lock_rois[0][3])
          is_inst_in = false;
      }
      bool has_this_cat = false;
      if (this->_categories_filter.size() == 0)  has_this_cat = true;
      for (int j=0; j<(int)this->_categories_filter.size(); j++)
      {
        if (tgts_json["targets"][i]["category_name"] == this->_categories_filter[j])  has_this_cat = true;
      }
      if (has_this_cat && is_inst_in)
      {
        found = true;
        if (!this->_locked)
        {
          this->_unlock_count = 0;
          if (this->_lock_count > 0)
          {
            if (sqrt((this->_unlock_last_cx-cx)*(this->_unlock_last_cx-cx) + (this->_unlock_last_cy-cy)*(this->_unlock_last_cy-cy)) < 0.2f * std::max(img_h, img_w))
              this->_lock_count ++;
          }
          else
          {
            this->_lock_count ++;
          }
          this->_unlock_last_cx = cx;
          this->_unlock_last_cy = cy;
          if (this->_lock_count >= this->_lock_thres)
          {
            this->_lock_count = 0;
            this->_locked = true;
            this->_unlock_last_cx = 0;
            this->_unlock_last_cy = 0;
            if (this->_lock_scale_count < 0)
              this->_lock_scale_count = 0;
          }
        }
        if (this->_locked)
        {
          this->_lock_rois.clear();
          std::vector<int> roi;
          roi.push_back((int)round(cx - wh / 2.f));
          roi.push_back((int)round(cy - wh / 2.f));
          roi.push_back((int)round(cx + wh / 2.f));
          roi.push_back((int)round(cy + wh / 2.f));
          
          if (roi[0] < 0)  { roi[2] += (-roi[0]); roi[0] = 0; }
          if (roi[1] < 0)  { roi[3] += (-roi[1]); roi[1] = 0; }
          if (roi[2] >= img_w)  { roi[0] -= (roi[2] - img_w + 1); roi[2] = img_w - 1; }
          if (roi[3] >= img_h)  { roi[1] -= (roi[3] - img_h + 1); roi[3] = img_h - 1; }

          if (roi[0] < 0)  { roi[0] = 0; }
          if (roi[1] < 0)  { roi[1] = 0; }

          this->_lock_rois.push_back(roi);

          if (roi[2] - roi[0] <= 0 || roi[3] - roi[1] <= 0)
          {
            this->_locked = false;
          }
          else
          {
            outputing = true;
            if (this->_last_cx >= 0.f && this->_last_cy >= 0.f)
            {
              float dist = sqrt((this->_last_cx-cx)*(this->_last_cx-cx) + (this->_last_cy-cy)*(this->_last_cy-cy));
              if (dist < _min_dist)
              {
                _min_dist = dist;
                _min_dist_roi = this->_lock_rois[0];
                _min_dist_tgt = tgts_json["targets"][i];
                _min_dist_cx = cx;
                _min_dist_cy = cy;
                _min_dist_cnt ++;
              }
            }
            else
            {
              _min_dist_roi = this->_lock_rois[0];
              _min_dist_tgt = tgts_json["targets"][i];
              _min_dist_cx = cx;
              _min_dist_cy = cy;
              _min_dist_cnt ++;
              break;
            }
          }
        }
      }
    }


    if (_min_dist_cnt > 0)
    {
      tgts_json["targets"].clear();
      tgts_json["targets"].push_back(_min_dist_tgt);
      this->_lock_rois.clear();
      this->_lock_rois.push_back(_min_dist_roi);
      this->_last_cx = _min_dist_cx;
      this->_last_cy = _min_dist_cy;
    }

    if (!found)
    {
      this->_lock_count = 0;
      this->_unlock_count ++;
      if (this->_unlock_count >= this->_unlock_thres)
      {
        this->_unlock_count = 0;
        this->_lock_scale_count = -1;
        this->_locked = false;
        this->_last_cx = -1.f;
        this->_last_cy = -1.f;
        this->_lock_rois.clear();
        tgts_json["rois"].clear();
        if (this->_use_square_region)
        {
          if (img_w > img_h)
          {
            tgts_json["rois"].push_back({(img_w-img_h) / 2, 0, img_h, img_h});
            this->_lock_rois.push_back({(img_w-img_h) / 2, 0, (img_w-img_h) / 2 + img_h, img_h});
          }
          else
          {
            tgts_json["rois"].push_back({0, (img_h-img_w) / 2, img_w, img_w});
            this->_lock_rois.push_back({0, (img_h-img_w) / 2, img_w, (img_h-img_w) / 2 + img_w});
          }
        }
        else
        {
          ;
        }
      }
    }

    if (this->_locked)
    {
      this->_lock_scale_count ++;
    }

    if (!outputing && !this->_keep_unlocked)
    {
      tgts_json["targets"].clear();
    }


    image_json["spirecv_msgs::2DTargets"] = tgts_json;
    _res_pub.publish(tgts_json);
    _res_img_pub.publish(image_json);


    if (this->_locked && this->_lock_rois.size() > 0)
    {
      nlohmann::json roi_ctl_msg = sms::def_msg("std_msgs::Null");
      roi_ctl_msg["rois"] = nlohmann::json::array();
      for (int j=0; j<this->_lock_rois.size(); j++)
      {
        nlohmann::json roi_json = nlohmann::json::array();
        roi_json.push_back(this->_lock_rois[j][0]);
        roi_json.push_back(this->_lock_rois[j][1]);
        roi_json.push_back(this->_lock_rois[j][2]);
        roi_json.push_back(this->_lock_rois[j][3]);
        roi_ctl_msg["rois"].push_back(roi_json);
      }
      roi_ctl_msg["pause"] = false;
      roi_ctl_msg["input_size"] = 640;
      _roi_det_ctl_pub.publish(roi_ctl_msg);

      nlohmann::json full_ctl_msg = sms::def_msg("std_msgs::Null");
      full_ctl_msg["pause"] = true;
      full_ctl_msg["input_size"] = 1280;
      _full_det_ctl_pub.publish(full_ctl_msg);
    }
    else
    {
      nlohmann::json roi_ctl_msg = sms::def_msg("std_msgs::Null");
      roi_ctl_msg["pause"] = true;
      roi_ctl_msg["input_size"] = 640;
      _roi_det_ctl_pub.publish(roi_ctl_msg);

      nlohmann::json full_ctl_msg = sms::def_msg("std_msgs::Null");
      full_ctl_msg["pause"] = false;
      full_ctl_msg["input_size"] = 1280;
      _full_det_ctl_pub.publish(full_ctl_msg);
    }

  }
}



}
